#include "gtest/gtest.h"
#include <mockcpp/mokc.h>
#include <mockcpp/mockcpp.hpp>
 
#include <vector>
#include <iostream>
 
#include "topoinfo_struct.h"
#include "log.h"
#include "checker.h"
#include "virtual_topo.h"
#include "testcase_utils.h"
#include "coll_reduce_scatter_v_executor.h"
 
using namespace hccl;

class ReduceScatterVTestFor310P : public testing::Test {
protected:
    static void SetUpTestCase()
    {
        std::cout << "ReduceScatterVTestFor310P set up." << std::endl;
    }
 
    static void TearDownTestCase()
    {
        std::cout << "ReduceScatterVTestFor310P tear down." << std::endl;
    }
    virtual void SetUp()
    {
        const ::testing::TestInfo* const test_info = ::testing::UnitTest::GetInstance()->current_test_info();
        std::string caseName = "analysis_result_" + std::string(test_info->test_case_name()) + "_" + std::string(test_info->name());
        Checker::SetDumpFileName(caseName);
    }
 
    virtual void TearDown()
    {
        Checker::SetDumpFileName("analysis_result");
        GlobalMockObject::verify();
        // 这边每个case执行完成需要清理所有的环境变量，如果有新增的环境变量，需要在这个函数中进行清理
        ClearHcclEnv();
    }
};
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 4);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts;
    vector<u64> displs;
    u64 displacement = 0;
    const u64 count = 100;
    for (u32 i = 0; i < ranktable.rankNum; i++) {
        counts.emplace_back(count);
        displs.emplace_back(displacement);
        displacement += count;
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 4);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {100, 200, 300, 400};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_2)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 4);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {400, 300, 200, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_3)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 2);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {400, 300};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_4)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 5);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {100, 200, 300, 200, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_large_size)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 4);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {40000000 , 30000000, 20000000, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_tbe)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 4);
    gen.GenRankTable(topoMeta, ranktable);

    vector<u64> counts {400, 300, 200, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT16;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_6)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 6);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {100, 200, 300, 300, 200, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_7)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 5);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {100, 200, 300, 400, 300, 200, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_8)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 5);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {100, 200, 300, 400, 400, 300, 200, 100};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_FP32;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}
 
TEST_F(ReduceScatterVTestFor310P, reduce_scatter_v_310P3_opbase_ReduceScatterRing_varying_size_containes_zero)
{
    RankTable_t ranktable;
    RankTable_For_LLT gen;
    TopoMeta topoMeta;
    gen.GenTopoMeta(topoMeta, 1, 1, 4);
    gen.GenRankTable(topoMeta, ranktable);
 
    vector<u64> counts {400, 0, 200, 0};
    vector<u64> displs {0};
    for (auto i = 1; i < counts.size(); ++i) {
            displs.emplace_back(displs[i-1] + counts[i-1]);
    }
 
    TestOpParam testOpParam;
    testOpParam.opType = HcclCMDType::HCCL_CMD_REDUCE_SCATTER_V;
    testOpParam.tag = "ReduceScatterV";
    testOpParam.opMode = OPBASE;
    testOpParam.devtype = DevType::DEV_TYPE_310P3;
    testOpParam.VDataDes.counts = counts;
    testOpParam.VDataDes.displs = displs;
    testOpParam.VDataDes.dataType = HcclDataType::HCCL_DATA_TYPE_INT16;
    testOpParam.reduceType = HcclReduceOp::HCCL_REDUCE_SUM;
    testOpParam.algName = "ReduceScatterVFor310PRing";
 
    Checker checker;
    HcclResult ret;
    ret = checker.Check(testOpParam, ranktable, topoMeta);
    EXPECT_EQ(ret, HcclResult::HCCL_SUCCESS);
}